# -*- coding: utf-8 -*-

from spectrogram import logmelspectrogram
import numpy as np
from joblib import Parallel, delayed
import librosa
import soundfile as sf
import os
from glob import glob
from tqdm import tqdm
import random
import json
import resampy
import pyworld as pw
# from fairseq.models.wav2vec import Wav2Vec2Model
import torch
import fairseq

def get_mean_std():
    
    data_root = '/root/VCTK/Hubert/train/mels'
    save_root = '/opt/data/private/VC/VCTK/Hubert'

    # 读取train文件夹的npy文件
    all_mels = glob(f'{data_root}/*/*mic1.npy')

    # 聚合之后做均值 方差
    mels = []
    for wav_path in tqdm(all_mels):
        mel=np.load(wav_path)
        mels.append(mel)

    mels = np.concatenate(mels, 0)
    mean = np.mean(mels, 0)
    std = np.std(mels, 0)

    # 存储起来
    mel_stats = np.concatenate([mean.reshape(1,-1), std.reshape(1,-1)], 0)
    np.save(f'{save_root}/Hubert_mel_stats.npy', mel_stats)

# mel=np.load('/opt/data/private/VC/VCTK/Hubert_norlize/Hubert_mel_stats.npy')
# print(mel.shape)

def normalize():
    mel_status=np.load('/opt/data/private/VC/VCTK/Hubert_norlize/Hubert_mel_stats.npy')
    mean=mel_status[0].reshape(1,-1)
    std=mel_status[1].reshape(1,-1)
    
    
    data_root = '/opt/data/private/VC/VCTK/Hubert/test/mels'
    save_root = '/opt/data/private/VC/VCTK/Hubert_norlize/test/mels'

    # 读取train文件夹的npy文件
    all_mels = glob(f'{data_root}/*/*mic1.npy')

    # 聚合之后做均值 方差
    mels = []
    for wav_path in tqdm(all_mels):
        mel=np.load(wav_path)
        mel = (mel - mean) / (std + 1e-8)
        
        file_name = wav_path.split('/')[-1]
        save_path=os.path.join(save_root,file_name)
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        np.save(save_path, mel)
        
get_mean_std()